package com.bigdata.iotdb;

import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.ParallelSourceFunction;
import org.apache.iotdb.jdbc.IoTDBSQLException;

import java.sql.*;
import java.util.ArrayList;
import java.util.List;

/**
 * 利用jdbc读取iotdb数据库，
 * 并将宽表信息构建为flink数据源
 */
public class ReadFromIoTDB {

    private static String sql;
    private static String host;
    public static void main(String[] args) throws Exception {
        sql = args[0];
        host = args[1];
        final StreamExecutionEnvironment senv = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStreamSource<FormData> stream = senv.addSource(new iotdbSource()).setParallelism(2);
        stream.print().setParallelism(1);
        senv.execute("ReadFromIoTDB");
    }

    private static class iotdbSource implements ParallelSourceFunction<FormData> {
        boolean isRunning = true;

        @Override
        public void run(SourceContext<FormData> ctx) throws Exception {
            Class.forName("org.apache.iotdb.jdbc.IoTDBDriver");
            while(isRunning){
                List<String> labels = new ArrayList<>();
                List<List<String>> data = new ArrayList<>();
                try(Connection connection = DriverManager.getConnection(
                        "jdbc:iotdb://" + host + ":6667/", "root", "root");
                    Statement statement = connection.createStatement()) {
                    ResultSet resultSet = statement.executeQuery(sql);
                    ResultSetMetaData metaData = resultSet.getMetaData();
                    int columnCount = metaData.getColumnCount();
                    for (int i = 0; i < columnCount; i++) {
                        labels.add(metaData.getColumnLabel(i + 1));
                    }
                    while (resultSet.next()) {
                        List<String> rowdata = new ArrayList<>();
                        for (int i = 1; ; i++) {
                            rowdata.add(resultSet.getString(i));
                            if (i >= columnCount) {
                                data.add(rowdata);
                                break;
                            }
                        }
                    }
                }catch (IoTDBSQLException e){
                    System.out.println(e.getMessage());
                }
                FormData formData = new FormData();
                formData.setLabels(labels);
                formData.setData(data);
                ctx.collect(formData);
                Thread.sleep(1000);
            }
        }

        @Override
        public void cancel() {
            isRunning = false;
        }
    }
}
